"""LangGraph 多 Agent 协作 - 官方 Supervisor 模式 (v0.6.8)"""
import os
from typing import Annotated
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_core.tools import tool

from langgraph.graph import StateGraph, END, MessagesState, add_messages
from langgraph.prebuilt import create_react_agent

# 加载环境变量
load_dotenv()

# 配置
LLM_MODEL = os.getenv("LLM_MODEL", "gemini-2.5-flash-preview-05-20")
LLM_TEMPERATURE = float(os.getenv("LLM_TEMPERATURE", "0.3"))
LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "65535"))
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_API_BASE = os.getenv("OPENAI_API_BASE")

print(f"[CONFIG] Model: {LLM_MODEL}")

# 初始化模型
model = ChatOpenAI(
    model=LLM_MODEL,
    temperature=LLM_TEMPERATURE,
    max_tokens=LLM_MAX_TOKENS,
    api_key=OPENAI_API_KEY,
    base_url=OPENAI_API_BASE,
    streaming=True
)

print(f"[INIT] Model initialized")

# ========== 定义 Agent 工具 ==========

@tool
def math_agent(query: str) -> str:
    """
    专业数学助手。用于：
    - 数学计算（加减乘除、幂运算等）
    - 解方程和数学问题解答
    - 提供详细的计算步骤

    Args:
        query: 用户的数学问题

    Returns:
        数学问题的答案和计算过程
    """
    print(f"\n[MATH_AGENT] Processing: {query[:50]}...")

    system_prompt = """你是专业的数学助手。
请简洁清晰地解答数学问题，提供计算步骤。"""

    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=query)
    ]

    response = model.invoke(messages)
    print(f"[MATH_AGENT] Completed: {response.content[:80]}...\n")

    return response.content


@tool
def poet_agent(query: str) -> str:
    """
    富有创意的诗人。用于：
    - 诗歌创作（现代诗、古诗、各种风格）
    - 文学表达和情感抒发
    - 优美的文字创作

    Args:
        query: 用户的诗歌创作需求或主题

    Returns:
        创作的诗歌内容
    """
    print(f"\n[POET_AGENT] Creating: {query[:50]}...")

    system_prompt = """你是富有创意的诗人。
请用优美诗意的语言回应，创作打动人心的诗歌。"""

    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=query)
    ]

    response = model.invoke(messages)
    print(f"[POET_AGENT] Completed: {response.content[:80]}...\n")

    return response.content


# 工具列表
tools = [math_agent, poet_agent]

# ========== 定义状态 ==========
from typing import TypedDict

class AgentState(TypedDict):
    """Agent 状态定义 - 兼容 LangGraph 0.6.8"""
    messages: Annotated[list, add_messages]
    remaining_steps: int  # LangGraph 0.6.8 要求的字段

# ========== Supervisor 节点（使用 create_react_agent）==========

# Supervisor 系统提示 - 通过包装实现
def create_supervisor_agent():
    """创建带系统提示的 Supervisor Agent"""

    system_message = SystemMessage(content="""你是一个智能任务调度主管(Supervisor)。

你的职责:
1. 理解用户的需求
2. 选择最合适的专家工具来处理任务
3. 将任务委派给专家执行

可用的专家工具:
- math_agent: 数学计算和问题解答专家
- poet_agent: 诗歌创作和文学表达专家

决策规则:
- 如果用户问题涉及数学计算、数字运算、解题 → 使用 math_agent
- 如果用户问题涉及诗歌创作、文学表达、情感抒发 → 使用 poet_agent

请分析用户问题，选择合适的工具，并将用户的完整问题传递给工具。
""")

    # 创建基础 react agent
    base_agent = create_react_agent(
        model,
        tools=tools,
        state_schema=AgentState  # 0.6.x 版本使用 state_schema
    )

    print(f"[SUPERVISOR] React agent created with {len(tools)} tools")

    # 包装 agent 以添加系统提示
    def supervisor_with_system_prompt(state: AgentState):
        """带系统提示的 Supervisor"""
        # 在消息前添加系统提示
        messages = state.get("messages", [])

        # 检查是否已有系统消息
        has_system = any(isinstance(msg, SystemMessage) for msg in messages)

        if not has_system:
            # 添加系统提示
            modified_state = {
                "messages": [system_message] + messages
            }
        else:
            modified_state = state

        # 调用 base agent
        return base_agent.invoke(modified_state)

    return supervisor_with_system_prompt


# ========== 构建 Graph ==========
def create_workflow():
    """创建 Workflow"""
    print(f"\n[BUILD] Creating Supervisor workflow...")

    # 创建 Supervisor
    supervisor = create_supervisor_agent()

    # 创建 Graph
    workflow = StateGraph(AgentState)

    # 添加 Supervisor 节点
    workflow.add_node("supervisor", supervisor)

    # 设置入口
    workflow.set_entry_point("supervisor")

    # Supervisor 执行完直接结束
    workflow.add_edge("supervisor", END)

    print(f"[BUILD] Workflow created (Official Supervisor Pattern)\n")

    return workflow.compile()


# 导出 graph
graph = create_workflow()
print(f"[EXPORT] Graph ready with LangGraph Supervisor pattern\n")


# ========== 测试代码 ==========
if __name__ == "__main__":
    print("=" * 80)
    print("[TEST] Testing Official Supervisor Pattern")
    print("=" * 80)

    # 测试 1: 数学问题
    print("\n" + "=" * 80)
    print("[TEST 1] Math question: 计算 25 * 4 + 10")
    print("=" * 80)

    try:
        result = graph.invoke({
            "messages": [HumanMessage(content="计算 25 * 4 + 10")]
        })

        print(f"\n[RESULT] Messages in state: {len(result['messages'])}")
        print(f"[FINAL] Response:")
        # 找到最后的非工具调用消息
        for msg in reversed(result["messages"]):
            if isinstance(msg, AIMessage):
                # 跳过工具调用消息
                if not getattr(msg, 'tool_calls', None):
                    print(f"   {msg.content}")
                    break
    except Exception as e:
        print(f"[ERROR] Math test failed: {e}")

    # 测试 2: 诗歌创作
    print("\n" + "=" * 80)
    print("[TEST 2] Poetry creation: 写一首关于春天的短诗")
    print("=" * 80)

    try:
        result = graph.invoke({
            "messages": [HumanMessage(content="写一首关于春天的短诗")]
        })

        print(f"\n[RESULT] Messages in state: {len(result['messages'])}")
        print(f"[FINAL] Response:")
        for msg in reversed(result["messages"]):
            if isinstance(msg, AIMessage):
                if not getattr(msg, 'tool_calls', None):
                    print(f"   {msg.content}")
                    break
    except Exception as e:
        print(f"[ERROR] Poetry test failed: {e}")

    # 测试 3: 混合任务
    print("\n" + "=" * 80)
    print("[TEST 3] Mixed task: 计算123+456，然后为结果写首诗")
    print("=" * 80)

    try:
        result = graph.invoke({
            "messages": [HumanMessage(content="先计算123+456，然后为这个结果写一首诗")]
        })

        print(f"\n[RESULT] Messages in state: {len(result['messages'])}")
        print(f"[FINAL] All responses:")
        for i, msg in enumerate(result["messages"]):
            if isinstance(msg, AIMessage) and not getattr(msg, 'tool_calls', None):
                print(f"   [{i}] {msg.content[:100]}...")
    except Exception as e:
        print(f"[ERROR] Mixed task test failed: {e}")

    print("\n" + "=" * 80)
    print("[TEST] All tests completed")
    print("=" * 80)